#include <stdio.h>
#include "blas.h"
#include "network.h"

network *make_network(int n)
{
    network *net = calloc(1, sizeof(network));
    net->n = n;
    net->layers = calloc(net->n, sizeof(layer));
    net->cost = calloc(1, sizeof(float));
    net->seen = calloc(1, sizeof(size_t));
    net->subdivisions = 1;
    return net;
}

void calc_network_cost(network *netp)
{
    network net = *netp;
    int i;
    float sum = 0;
    int count = 0;
    for (i = 0; i < net.n; i++) {
        if (net.layers[i].cost) {
            sum += net.layers[i].cost[0];
            ++count;
        }
    }
    *net.cost = sum / count;
}
void forward_network(network *netp)
{
    int i;
    network net = *netp;

    for (i = 0; i < net.n; i++) {
        layer l = net.layers[i]; 
        net.index = i;
        if (l.delta) {
            fill_cpu(l.outputs * l.batch, 0, l.delta, 1);
        }
        l.forward(l, net);
        net.input = l.output;
        if (l.truth) {
            net.truth = l.output;
        }
    }
    calc_network_cost(netp);
}

void backward_network(network *netp)
{
    network net = *netp;
    int i;
    network orig = net;
    for (i = net.n - 1; i > 0; i--) {
        layer l = net.layers[i];
        if (l.stopbackward) break;
        if (i == 1) {
            net = orig;
            if (net.delta == NULL) {
                layer prev = net.layers[0];
                net.input = prev.output;
                net.delta = prev.delta;
            }
        } else {
            layer prev = net.layers[i - 1];
            net.input = prev.output;
            net.delta = prev.delta;
        }
        net.index = i;
        l.backward(l, net);
    }
}

void update_network(network *netp)
{
    network net = *netp;
    int i;
    update_args a = {0};
    a.batch = net.batch * net.subdivisions;
}

float train_network_datum(network *net)
{
    float error;
    int update_count;
    *net->seen += net->batch;
    net->train = 1;
    forward_network(net);
    backward_network(net);

    error = *net->cost;
    update_count = (*net->seen)/ net->batch;
    update_count = update_count % net->subdivisions;
    if (!update_count) {
        update_network(net);
    }
    return error;
}

void free_network(network *net)
{
    int i;
    for (i = 0; i < net->n; i++) {
        free_layer(net->layers[i]);
    }
    
    free(net->layers);

    if (net->cost) {
        free(net->cost);
        net->cost = NULL;
    }
    if (net->seen) {
        free(net->seen);
        net->seen = NULL;
    }
    free(net);
}